#include <dai/factorgraph.h>
#include <dai/bp.h>
#include <dai/hak.h>
#include <dai/lc.h>
#include <dai/jtree.h>
#include <dai/mf.h>
#include <dai/bbp.h>
#include <dai/alldai.h>
#include <vector>
#include <fstream>
#include <iostream>

/*
Inference algorithms:
Don't work on this example:
BP, BBP (can't clamp), MF, HAK (with other options tried)

Work on the example: HAK (with opts ("doubleloop",true) and ("clusters",HAK::Properties::ClustersType::MIN),
JTree (with opts ("updates","HUGIN"))
 */

#define __INF_ALG JTree
#define __UP_S "HUGIN"
/*
#define __INF_ALG HAK
#define __INF_ALG JTree
#define __UP_S "HUGIN"
*/

using namespace std;
using namespace dai;

class BPwrapper{
public:
  BPwrapper(){
    graph = FactorGraph();
    graph.ReadFromFile("./tempdai.fg");
    fstream dotOut("test.dot",fstream::out);
    graph.printDot(dotOut);
    clamped = false;
    PropertySet opts;
    size_t maxiter = 10000;
    Real tol = 1e-9;
    size_t verb = 0;
    //HAK::Properties::ClustersType clusters = HAK::Properties::ClustersType::MIN;
    opts.set("maxiter",maxiter);  // Maximum number of iterations
    opts.set("tol",tol);          // Tolerance for convergence
    opts.set("verbose",verb);
    opts.set("logdomain",true);
    opts.set("updates",string(__UP_S));
    /*opts.set("doubleloop",true);
      opts.set("clusters",clusters);*/
    //bp = BP(graph,opts("updates",string("SEQRND"))("logdomain",true));
    //bp = __INF_ALG(graph,opts("updates",string(__UP_S))("logdomain",true));
    bp = __INF_ALG(graph,opts);
    bp.init();
    bp.run();
    //cout << bp.beliefs() << endl;
  }
  void addVar(const int ind) {
    vlist.push_back(graph.var(ind));
  }
  void clamp(const int ind, const bool val) {
    //cout << "clamping node with id " << ind << " to value " << val << endl;
    if(clamped) {
      //cgraph = cgraph.clamped(ind,val);
      //graph.clamp(ind,val);
      bp.clamp(ind,val);
    }
    else {
      //cgraph = graph.clamped(ind,val);
      //bgraph = graph;
      bbp = bp;
      //graph.clamp(ind,val);
      bp.clamp(ind,val);
      clamped = true;
    }
  }
  void resetClamps() {
    clamped = false;
    //graph = bgraph;
    bp = bbp;
  }

  void run() {
    bp.run();
  }

  double getJoint(const int ind) {
    if (ind > ((2 << vlist.size()) - 1)) {
      return -1;
    }
    //bp.run();
    VarSet vset = VarSet(vlist.end()-1,vlist.end());
    
    return bp.belief(vset)[1-ind];
  }

  /*double getSJoint(const int id, const int ind) {
    if (ind > ((2 << vlist.size()) - 1))
      return -1;
    return bp.beliefV(id)[1-ind];
    }*/

  void resetVars() {
    vlist = vector<Var>();
  }
private:
  //BP bp;
  __INF_ALG bp;
  //BP bbp;
  __INF_ALG bbp;
  FactorGraph graph;
  FactorGraph cgraph;
  FactorGraph bgraph;
  vector<Var> vlist;
  bool clamped;
};

extern "C" {
  BPwrapper* newBP() {return new BPwrapper();}
  void addVar(BPwrapper* BP, const int ind) {BP->addVar(ind);}
  double getJoint(BPwrapper* BP, const int ind) {return BP->getJoint(ind);}
  //double getSJoint(BPwrapper* BP, const int id, const int ind) {return BP->getSJoint(id,ind);}
  void dealloc(BPwrapper* fpoint) {delete(fpoint);}
  void clamp(BPwrapper* BP, const int ind, const bool val) {BP->clamp(ind,val);}
  void resetClamps(BPwrapper* BP) {BP->resetClamps();}
  void run(BPwrapper* BP) {BP->run();}
  void resetVars(BPwrapper* BP) {BP->resetVars();}
}

